

<!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  
  <title>mindspore.ops.composite.base &mdash; MindSpore master documentation</title>
  

  
  <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />

  
  

  
  

  

  
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../../../" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/language_data.js"></script>
        <script async="async" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script>
    
    <script type="text/javascript" src="../../../../_static/js/theme.js"></script>

    
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../../../index.html" class="icon icon-home"> MindSpore
          

          
          </a>

          
            
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption"><span class="caption-text">MindSpore Python API</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.html">mindspore</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.common.initializer.html">mindspore.common.initializer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.communication.html">mindspore.communication</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.compression.html">mindspore.compression</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.context.html">mindspore.context</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.html">mindspore.dataset</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.audio.html">mindspore.dataset.audio</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.config.html">mindspore.dataset.config</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.text.html">mindspore.dataset.text</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.transforms.html">mindspore.dataset.transforms</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.vision.html">mindspore.dataset.vision</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.mindrecord.html">mindspore.mindrecord</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.html">mindspore.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.probability.html">mindspore.nn.probability</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.transformer.html">mindspore.nn.transformer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.numpy.html">mindspore.numpy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.ops.html">mindspore.ops</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.html">mindspore.parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.nn.html">mindspore.parallel.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.profiler.html">mindspore.profiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.scipy.html">mindspore.scipy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.train.html">mindspore.train</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.boost.html">mindspore.boost</a></li>
</ul>
<p class="caption"><span class="caption-text">MindSpore C++ API</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://www.mindspore.cn/lite/api/zh-CN/master/api_cpp/mindspore.html">MindSpore Lite↗</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">

      
      <nav class="wy-nav-top" aria-label="top navigation">
        
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">MindSpore</a>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        
          

















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
        
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
        
      <li>mindspore.ops.composite.base</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <h1>Source code for mindspore.ops.composite.base</h1><div class="highlight"><pre>
<span></span><span class="c1"># This is the Python adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).</span>
<span class="c1">#</span>
<span class="c1"># Copyright 2020-2021 Huawei Technologies Co., Ltd</span>
<span class="c1">#</span>
<span class="c1"># Licensed under the Apache License, Version 2.0 (the &quot;License&quot;);</span>
<span class="c1"># you may not use this file except in compliance with the License.</span>
<span class="c1"># You may obtain a copy of the License at</span>
<span class="c1">#</span>
<span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span>
<span class="c1">#</span>
<span class="c1"># Unless required by applicable law or agreed to in writing, software</span>
<span class="c1"># distributed under the License is distributed on an &quot;AS IS&quot; BASIS,</span>
<span class="c1"># WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.</span>
<span class="c1"># See the License for the specific language governing permissions and</span>
<span class="c1"># limitations under the License.</span>
<span class="c1"># ============================================================================</span>

<span class="sd">&quot;&quot;&quot;Basic composite operations.&quot;&quot;&quot;</span>
<span class="kn">from</span> <span class="nn">functools</span> <span class="kn">import</span> <span class="n">partial</span>
<span class="kn">from</span> <span class="nn">types</span> <span class="kn">import</span> <span class="n">FunctionType</span>

<span class="kn">from</span> <span class="nn">mindspore</span> <span class="kn">import</span> <span class="n">context</span>
<span class="kn">from</span> <span class="nn">..._c_expression</span> <span class="kn">import</span> <span class="n">GradOperation_</span><span class="p">,</span> <span class="n">HyperMap_</span><span class="p">,</span> <span class="n">Map_</span><span class="p">,</span> <span class="n">MultitypeFuncGraph_</span><span class="p">,</span> <span class="n">Tail_</span><span class="p">,</span> <span class="n">Shard_</span><span class="p">,</span> \
    <span class="n">TupleAdd_</span><span class="p">,</span> <span class="n">TupleSlice_</span><span class="p">,</span> <span class="n">UnpackCall_</span><span class="p">,</span> <span class="n">ZipOperation_</span><span class="p">,</span> <span class="n">ListAppend_</span><span class="p">,</span> <span class="n">TupleGetItemTensor_</span><span class="p">,</span> <span class="n">ListInsert_</span>
<span class="kn">from</span> <span class="nn">...common</span> <span class="kn">import</span> <span class="n">dtype</span> <span class="k">as</span> <span class="n">mstype</span>
<span class="kn">from</span> <span class="nn">...common.api</span> <span class="kn">import</span> <span class="n">ms_function</span><span class="p">,</span> <span class="n">_pynative_executor</span><span class="p">,</span> <span class="n">_wrap_func</span>
<span class="kn">from</span> <span class="nn">..primitive</span> <span class="kn">import</span> <span class="n">Primitive</span>
<span class="kn">from</span> <span class="nn">..operations</span> <span class="kn">import</span> <span class="n">_grad_ops</span>
<span class="kn">from</span> <span class="nn">..</span> <span class="kn">import</span> <span class="n">operations</span> <span class="k">as</span> <span class="n">P</span>
<span class="kn">from</span> <span class="nn">..</span> <span class="kn">import</span> <span class="n">signature</span> <span class="k">as</span> <span class="n">sig</span>

<span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span><span class="n">TupleAdd_</span><span class="p">,</span> <span class="n">TupleSlice_</span><span class="p">,</span> <span class="n">UnpackCall_</span><span class="p">,</span> <span class="n">TupleGetItemTensor_</span><span class="p">]</span>


<span class="k">def</span> <span class="nf">add_flags</span><span class="p">(</span><span class="n">fn</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="o">**</span><span class="n">flags</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A decorator that adds a flag to the function.</span>

<span class="sd">    Note:</span>
<span class="sd">        Only supports bool value.</span>

<span class="sd">    Args:</span>
<span class="sd">        fn (Function): Function or cell to add flag. Default: None.</span>
<span class="sd">        flags (dict): Flags use kwargs. Default: None.</span>

<span class="sd">    Returns:</span>
<span class="sd">        Function, the function with added flags.</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; net = Net();</span>
<span class="sd">        &gt;&gt;&gt; net = add_flags(net, predit=True)</span>
<span class="sd">        &gt;&gt;&gt; print(hasattr(net, &#39;_mindspore_flags&#39;))</span>
<span class="sd">        True</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="nf">deco</span><span class="p">(</span><span class="n">fn</span><span class="p">):</span>
        <span class="c1"># need set the attr and access on c++</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="s2">&quot;_mindspore_flags&quot;</span><span class="p">):</span>
            <span class="n">fn</span><span class="o">.</span><span class="n">_mindspore_flags</span> <span class="o">=</span> <span class="p">{}</span>

        <span class="n">fn</span><span class="o">.</span><span class="n">_mindspore_flags</span><span class="o">.</span><span class="n">update</span><span class="p">({</span><span class="o">**</span><span class="n">flags</span><span class="p">})</span>
        <span class="k">return</span> <span class="n">fn</span>
    <span class="n">ret</span> <span class="o">=</span> <span class="n">deco</span>
    <span class="k">if</span> <span class="n">fn</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="n">deco</span><span class="p">(</span><span class="n">fn</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">ret</span>


<div class="viewcode-block" id="core"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.core.html#mindspore.ops.core">[docs]</a><span class="k">def</span> <span class="nf">core</span><span class="p">(</span><span class="n">fn</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="o">**</span><span class="n">flags</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A decorator that adds a flag to the function.</span>

<span class="sd">    By default, the function is marked as True, enabling to use this decorator to</span>
<span class="sd">    set flag to a graph.</span>

<span class="sd">    Args:</span>
<span class="sd">        fn (Function): Function to add flag. Default: None.</span>
<span class="sd">        flags (dict): The following flags can be set core, which indicates that this is a core function or</span>
<span class="sd">                      other flag. Default: None.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; net = Net()</span>
<span class="sd">        &gt;&gt;&gt; net = core(net, predit=True)</span>
<span class="sd">        &gt;&gt;&gt; print(hasattr(net, &#39;_mindspore_flags&#39;))</span>
<span class="sd">        True</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># need set the attr and access on c++</span>

    <span class="k">def</span> <span class="nf">deco</span><span class="p">(</span><span class="n">fn</span><span class="p">):</span>
        <span class="n">fn</span><span class="o">.</span><span class="n">_mindspore_flags</span> <span class="o">=</span> <span class="p">{</span>
            <span class="s1">&#39;core&#39;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
            <span class="o">**</span><span class="n">flags</span><span class="p">,</span>
        <span class="p">}</span>
        <span class="k">return</span> <span class="n">fn</span>

    <span class="k">if</span> <span class="n">fn</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="n">deco</span><span class="p">(</span><span class="n">fn</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="n">deco</span>
    <span class="k">return</span> <span class="n">ret</span></div>


<div class="viewcode-block" id="GradOperation"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.GradOperation.html#mindspore.ops.GradOperation">[docs]</a><span class="k">class</span> <span class="nc">GradOperation</span><span class="p">(</span><span class="n">GradOperation_</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A higher-order function which is used to generate the gradient function for the input function.</span>

<span class="sd">    The gradient function generated by `GradOperation` higher-order function can be customized by</span>
<span class="sd">    construction arguments.</span>

<span class="sd">    Given an input function `net = Net()` that takes `x` and `y` as inputs, and has a parameter `z`,</span>
<span class="sd">    see `Net` in Examples.</span>


<span class="sd">    To generate a gradient function that returns gradients with respect to the first input</span>
<span class="sd">    (see `GradNetWrtX` in Examples).</span>

<span class="sd">    1. Construct a `GradOperation` higher-order function with default arguments:</span>
<span class="sd">       `grad_op = GradOperation()`.</span>

<span class="sd">    2. Call it with input function as argument to get the gradient function: `gradient_function = grad_op(net)`.</span>

<span class="sd">    3. Call the gradient function with input function&#39;s inputs to get the gradients with respect to the first input:</span>
<span class="sd">       `grad_op(net)(x, y)`.</span>


<span class="sd">    To generate a gradient function that returns gradients with respect to all inputs (see `GradNetWrtXY` in Examples).</span>

<span class="sd">    1. Construct a `GradOperation` higher-order function with `get_all=True` which</span>
<span class="sd">       indicates getting gradients with respect to all inputs, they are `x` and `y` in example function `Net()`:</span>
<span class="sd">       `grad_op = GradOperation(get_all=True)`.</span>

<span class="sd">    2. Call it with input function as argument to get the gradient function: `gradient_function = grad_op(net)`.</span>

<span class="sd">    3. Call the gradient function with input function&#39;s inputs to get the gradients with respect to all inputs:</span>
<span class="sd">       `gradient_function(x, y)`.</span>

<span class="sd">    To generate a gradient function that returns gradients with respect to given parameters</span>
<span class="sd">    (see `GradNetWithWrtParams` in Examples).</span>

<span class="sd">    1. Construct a `GradOperation` higher-order function with `get_by_list=True`:</span>
<span class="sd">       `grad_op = GradOperation(get_by_list=True)`.</span>

<span class="sd">    2. Construct a `ParameterTuple` that will be passed to the input function when constructing</span>
<span class="sd">       `GradOperation` higher-order function, it will be used as a parameter filter that determine</span>
<span class="sd">       which gradient to return: `params = ParameterTuple(net.trainable_params())`.</span>

<span class="sd">    3. Call it with input function and `params` as arguments to get the gradient function:</span>
<span class="sd">       `gradient_function = grad_op(net, params)`.</span>

<span class="sd">    4. Call the gradient function with input function&#39;s inputs to get the gradients with</span>
<span class="sd">       respect to given parameters: `gradient_function(x, y)`.</span>

<span class="sd">    To generate a gradient function that returns gradients with respect to all inputs and given parameters</span>
<span class="sd">    in the format of ((dx, dy), (dz))(see `GradNetWrtInputsAndParams` in Examples).</span>

<span class="sd">    1. Construct a `GradOperation` higher-order function with `get_all=True` and `get_by_list=True`:</span>
<span class="sd">       `grad_op = GradOperation(get_all=True, get_by_list=True)`.</span>

<span class="sd">    2. Construct a `ParameterTuple` that will be passed along input function when constructing</span>
<span class="sd">       `GradOperation` higher-order function: `params = ParameterTuple(net.trainable_params())`.</span>

<span class="sd">    3. Call it with input function and `params` as arguments to get the gradient function:</span>
<span class="sd">       `gradient_function = grad_op(net, params)`.</span>

<span class="sd">    4. Call the gradient function with input function&#39;s inputs</span>
<span class="sd">       to get the gradients with respect to all inputs and given parameters: `gradient_function(x, y)`.</span>


<span class="sd">    We can configure the sensitivity(gradient with respect to output) by setting `sens_param` as True and</span>
<span class="sd">    passing an extra sensitivity input to the gradient function, the sensitivity input should has the</span>
<span class="sd">    same shape and type with input function&#39;s output(see `GradNetWrtXYWithSensParam` in Examples).</span>

<span class="sd">    1. Construct a `GradOperation` higher-order function with `get_all=True` and `sens_param=True`:</span>
<span class="sd">       `grad_op = GradOperation(get_all=True, sens_param=True)`.</span>

<span class="sd">    2. Define `grad_wrt_output` as `sens_param` which works as the gradient with respect to output:</span>
<span class="sd">       `grad_wrt_output = Tensor(np.ones([2, 2]).astype(np.float32))`.</span>

<span class="sd">    3. Call it with input function as argument to get the gradient function:</span>
<span class="sd">       `gradient_function = grad_op(net)`.</span>

<span class="sd">    4. Call the gradient function with input function&#39;s inputs and `sens_param` to</span>
<span class="sd">       get the gradients with respect to all inputs:</span>
<span class="sd">       `gradient_function(x, y, grad_wrt_output)`.</span>

<span class="sd">    Args:</span>
<span class="sd">        get_all (bool): If True, get all the gradients with respect to inputs. Default: False.</span>
<span class="sd">        get_by_list (bool): If True, get all the gradients with respect to Parameter variables.</span>
<span class="sd">            If get_all and get_by_list are both False, get the gradient with respect to first input.</span>
<span class="sd">            If get_all and get_by_list are both True, get the gradients with respect to inputs and Parameter variables</span>
<span class="sd">            at the same time in the form of ((gradients with respect to inputs),</span>
<span class="sd">            (gradients with respect to parameters)). Default: False.</span>
<span class="sd">        sens_param (bool): Whether to append sensitivity (gradient with respect to output) as input.</span>
<span class="sd">            If sens_param is False, a &#39;ones_like(outputs)&#39; sensitivity will be attached automatically.</span>
<span class="sd">            Default: False.</span>
<span class="sd">            If the sensor_param is True, a sensitivity (gradient with respect to output) needs to be transferred</span>
<span class="sd">            through the location parameter or key-value pair parameter. If the value is transferred through</span>
<span class="sd">            the key-value pair parameter, the key must be sens.</span>

<span class="sd">    Returns:</span>
<span class="sd">        The higher-order function which takes a function as argument and returns gradient function for it.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `get_all`, `get_by_list` or `sens_param` is not a bool.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import ParameterTuple</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.ops.composite import GradOperation</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.ops import operations as P</span>
<span class="sd">        &gt;&gt;&gt; class Net(nn.Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(Net, self).__init__()</span>
<span class="sd">        ...         self.matmul = P.MatMul()</span>
<span class="sd">        ...         self.z = Parameter(Tensor(np.array([1.0], np.float32)), name=&#39;z&#39;)</span>
<span class="sd">        ...     def construct(self, x, y):</span>
<span class="sd">        ...         x = x * self.z</span>
<span class="sd">        ...         out = self.matmul(x, y)</span>
<span class="sd">        ...         return out</span>
<span class="sd">        ...</span>
<span class="sd">        &gt;&gt;&gt; class GradNetWrtX(nn.Cell):</span>
<span class="sd">        ...     def __init__(self, net):</span>
<span class="sd">        ...         super(GradNetWrtX, self).__init__()</span>
<span class="sd">        ...         self.net = net</span>
<span class="sd">        ...         self.grad_op = GradOperation()</span>
<span class="sd">        ...     def construct(self, x, y):</span>
<span class="sd">        ...         gradient_function = self.grad_op(self.net)</span>
<span class="sd">        ...         return gradient_function(x, y)</span>
<span class="sd">        ...</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor([[0.5, 0.6, 0.4], [1.2, 1.3, 1.1]], dtype=mstype.float32)</span>
<span class="sd">        &gt;&gt;&gt; y = Tensor([[0.01, 0.3, 1.1], [0.1, 0.2, 1.3], [2.1, 1.2, 3.3]], dtype=mstype.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = GradNetWrtX(Net())(x, y)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[1.4100001 1.5999999 6.6      ]</span>
<span class="sd">         [1.4100001 1.5999999 6.6      ]]</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; class GradNetWrtXY(nn.Cell):</span>
<span class="sd">        ...     def __init__(self, net):</span>
<span class="sd">        ...         super(GradNetWrtXY, self).__init__()</span>
<span class="sd">        ...         self.net = net</span>
<span class="sd">        ...         self.grad_op = GradOperation(get_all=True)</span>
<span class="sd">        ...     def construct(self, x, y):</span>
<span class="sd">        ...         gradient_function = self.grad_op(self.net)</span>
<span class="sd">        ...         return gradient_function(x, y)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor([[0.8, 0.6, 0.2], [1.8, 1.3, 1.1]], dtype=mstype.float32)</span>
<span class="sd">        &gt;&gt;&gt; y = Tensor([[0.1, 3.3, 1.1], [1.1, 0.2, 1.4], [1.1, 2.2, 0.3]], dtype=mstype.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = GradNetWrtXY(Net())(x, y)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        (Tensor(shape=[2, 3], dtype=Float32, value=</span>
<span class="sd">        [[ 4.50000000e+00,  2.70000005e+00,  3.60000014e+00],</span>
<span class="sd">         [ 4.50000000e+00,  2.70000005e+00,  3.60000014e+00]]), Tensor(shape=[3, 3], dtype=Float32, value=</span>
<span class="sd">        [[ 2.59999990e+00,  2.59999990e+00,  2.59999990e+00],</span>
<span class="sd">         [ 1.89999998e+00,  1.89999998e+00,  1.89999998e+00],</span>
<span class="sd">         [ 1.30000007e+00,  1.30000007e+00,  1.30000007e+00]]))</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; class GradNetWrtXYWithSensParam(nn.Cell):</span>
<span class="sd">        ...     def __init__(self, net):</span>
<span class="sd">        ...         super(GradNetWrtXYWithSensParam, self).__init__()</span>
<span class="sd">        ...         self.net = net</span>
<span class="sd">        ...         self.grad_op = GradOperation(get_all=True, sens_param=True)</span>
<span class="sd">        ...         self.grad_wrt_output = Tensor([[0.1, 0.6, 0.2], [0.8, 1.3, 1.1]], dtype=mstype.float32)</span>
<span class="sd">        ...     def construct(self, x, y):</span>
<span class="sd">        ...         gradient_function = self.grad_op(self.net)</span>
<span class="sd">        ...         return gradient_function(x, y, self.grad_wrt_output)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor([[0.8, 0.6, 0.2], [1.8, 1.3, 1.1]], dtype=mstype.float32)</span>
<span class="sd">        &gt;&gt;&gt; y = Tensor([[0.11, 3.3, 1.1], [1.1, 0.2, 1.4], [1.1, 2.2, 0.3]], dtype=mstype.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = GradNetWrtXYWithSensParam(Net())(x, y)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        (Tensor(shape=[2, 3], dtype=Float32, value=</span>
<span class="sd">        [[ 2.21099997e+00,  5.09999990e-01,  1.49000001e+00],</span>
<span class="sd">         [ 5.58800030e+00,  2.68000007e+00,  4.07000017e+00]]), Tensor(shape=[3, 3], dtype=Float32, value=</span>
<span class="sd">        [[ 1.51999998e+00,  2.81999993e+00,  2.14000010e+00],</span>
<span class="sd">         [ 1.09999990e+00,  2.04999995e+00,  1.54999995e+00],</span>
<span class="sd">         [ 9.00000036e-01,  1.54999995e+00,  1.25000000e+00]]))</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; class GradNetWithWrtParams(nn.Cell):</span>
<span class="sd">        ...     def __init__(self, net):</span>
<span class="sd">        ...         super(GradNetWithWrtParams, self).__init__()</span>
<span class="sd">        ...         self.net = net</span>
<span class="sd">        ...         self.params = ParameterTuple(net.trainable_params())</span>
<span class="sd">        ...         self.grad_op = GradOperation(get_by_list=True)</span>
<span class="sd">        ...     def construct(self, x, y):</span>
<span class="sd">        ...         gradient_function = self.grad_op(self.net, self.params)</span>
<span class="sd">        ...         return gradient_function(x, y)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor([[0.8, 0.6, 0.2], [1.8, 1.3, 1.1]], dtype=mstype.float32)</span>
<span class="sd">        &gt;&gt;&gt; y = Tensor([[0.11, 3.3, 1.1], [1.1, 0.2, 1.4], [1.1, 2.2, 0.3]], dtype=mstype.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = GradNetWithWrtParams(Net())(x, y)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        (Tensor(shape=[1], dtype=Float32, value= [ 2.15359993e+01]),)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; class GradNetWrtInputsAndParams(nn.Cell):</span>
<span class="sd">        ...     def __init__(self, net):</span>
<span class="sd">        ...         super(GradNetWrtInputsAndParams, self).__init__()</span>
<span class="sd">        ...         self.net = net</span>
<span class="sd">        ...         self.params = ParameterTuple(net.trainable_params())</span>
<span class="sd">        ...         self.grad_op = GradOperation(get_all=True, get_by_list=True)</span>
<span class="sd">        ...     def construct(self, x, y):</span>
<span class="sd">        ...         gradient_function = self.grad_op(self.net, self.params)</span>
<span class="sd">        ...         return gradient_function(x, y)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor([[0.1, 0.6, 1.2], [0.5, 1.3, 0.1]], dtype=mstype.float32)</span>
<span class="sd">        &gt;&gt;&gt; y = Tensor([[0.12, 2.3, 1.1], [1.3, 0.2, 2.4], [0.1, 2.2, 0.3]], dtype=mstype.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = GradNetWrtInputsAndParams(Net())(x, y)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        ((Tensor(shape=[2, 3], dtype=Float32, value=</span>
<span class="sd">        [[ 3.51999998e+00,  3.90000010e+00,  2.59999990e+00],</span>
<span class="sd">         [ 3.51999998e+00,  3.90000010e+00,  2.59999990e+00]]), Tensor(shape=[3, 3], dtype=Float32, value=</span>
<span class="sd">        [[ 6.00000024e-01,  6.00000024e-01,  6.00000024e-01],</span>
<span class="sd">         [ 1.89999998e+00,  1.89999998e+00,  1.89999998e+00],</span>
<span class="sd">         [ 1.30000007e+00,  1.30000007e+00,  1.30000007e+00]])), (Tensor(shape=[1], dtype=Float32, value=</span>
<span class="sd">         [ 1.29020004e+01]),))</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">get_all</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">get_by_list</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">sens_param</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize GradOperation.&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">get_all</span><span class="p">,</span> <span class="nb">bool</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;GradOperation&#39;, the &#39;get_all&#39; should be bool, but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">get_all</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">get_by_list</span><span class="p">,</span> <span class="nb">bool</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;GradOperation&#39;, the &#39;get_by_list&#39; should be bool, &quot;</span>
                            <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">get_by_list</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">sens_param</span><span class="p">,</span> <span class="nb">bool</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;GradOperation&#39;, the &#39;sens_param&#39; should be bool, &quot;</span>
                            <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">sens_param</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">get_all</span> <span class="o">=</span> <span class="n">get_all</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">get_by_list</span> <span class="o">=</span> <span class="n">get_by_list</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sens_param</span> <span class="o">=</span> <span class="n">sens_param</span>
        <span class="n">GradOperation_</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;grad&#39;</span><span class="p">,</span> <span class="n">get_all</span><span class="p">,</span> <span class="n">get_by_list</span><span class="p">,</span> <span class="n">sens_param</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">grad_fn</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fn</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">pynative_</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">def</span> <span class="nf">_pynative_forward_run</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">grad</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">kwargs</span><span class="p">,</span> <span class="n">fn</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot; Pynative forward run to build grad graph. &quot;&quot;&quot;</span>
        <span class="n">new_kwargs</span> <span class="o">=</span> <span class="n">kwargs</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">sens_param</span><span class="p">:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="s1">&#39;sens&#39;</span> <span class="ow">in</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                <span class="n">args</span> <span class="o">=</span> <span class="n">args</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">new_kwargs</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
                <span class="n">new_kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s1">&#39;sens&#39;</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">FunctionType</span><span class="p">):</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">check_run</span><span class="p">(</span><span class="n">grad</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">new_kwargs</span><span class="p">):</span>
                <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">set_grad_flag</span><span class="p">(</span><span class="kc">True</span><span class="p">)</span>
                <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">new_graph</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">new_kwargs</span><span class="p">)</span>
                <span class="n">output</span> <span class="o">=</span> <span class="n">fn</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">new_kwargs</span><span class="p">)</span>
                <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">end_graph</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">output</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">new_kwargs</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="c1"># Check if fn have run already</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">check_run</span><span class="p">(</span><span class="n">grad</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">new_kwargs</span><span class="p">):</span>
                <span class="n">fn</span><span class="o">.</span><span class="n">set_grad</span><span class="p">()</span>
                <span class="n">fn</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">new_kwargs</span><span class="p">)</span>
                <span class="n">fn</span><span class="o">.</span><span class="n">set_grad</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="n">weights</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">grad_fn</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">fn</span> <span class="o">==</span> <span class="n">fn</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">grad_fn</span>
        <span class="n">grad_</span> <span class="o">=</span> <span class="n">GradOperation</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">get_all</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_by_list</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">sens_param</span><span class="p">)</span>
        <span class="c1"># If calling Grad in GRAPH_MODE or calling Grad in ms_function, do grad in GRAPH_MODE</span>
        <span class="c1"># If calling Grad in pure PYNATIVE_MODE do grad in PYNATIVE_MODE</span>
        <span class="c1">#   In pure PYNATIVE_MODE the out layer after_grad just used to set pynative flag for inner GradOperation.</span>
        <span class="c1">#   In PYNATIVE_MODE calling Grad from ms_function, use the out layer after_grad do grad in GRAPH_MODE.</span>
        <span class="k">if</span> <span class="n">context</span><span class="o">.</span><span class="n">get_context</span><span class="p">(</span><span class="s2">&quot;mode&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="n">context</span><span class="o">.</span><span class="n">GRAPH_MODE</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_by_list</span><span class="p">:</span>
                <span class="nd">@ms_function</span>
                <span class="k">def</span> <span class="nf">after_grad</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">):</span>
                    <span class="k">return</span> <span class="n">grad_</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">weights</span><span class="p">)(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="nd">@ms_function</span>
                <span class="k">def</span> <span class="nf">after_grad</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">):</span>
                    <span class="k">return</span> <span class="n">grad_</span><span class="p">(</span><span class="n">fn</span><span class="p">)(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">pynative_</span><span class="p">:</span>
            <span class="nd">@_wrap_func</span>
            <span class="k">def</span> <span class="nf">after_grad</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                <span class="k">if</span> <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">check_graph</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                    <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;Another grad step is running&quot;</span><span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">_pynative_forward_run</span><span class="p">(</span><span class="n">grad_</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">kwargs</span><span class="p">,</span> <span class="n">fn</span><span class="p">)</span>
                <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">grad</span><span class="p">(</span><span class="n">grad_</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="n">weights</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,),</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                <span class="n">out</span> <span class="o">=</span> <span class="n">_pynative_executor</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">clear_grad</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                <span class="k">return</span> <span class="n">out</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">grad_</span><span class="o">.</span><span class="n">pynative_</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="c1"># after_grad of this branch can&#39;t use @ms_function, just directly call grad_</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_by_list</span><span class="p">:</span>
                <span class="k">def</span> <span class="nf">after_grad</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                    <span class="k">return</span> <span class="n">grad_</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">weights</span><span class="p">)(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">def</span> <span class="nf">after_grad</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                    <span class="k">return</span> <span class="n">grad_</span><span class="p">(</span><span class="n">fn</span><span class="p">)(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">grad_fn</span> <span class="o">=</span> <span class="n">after_grad</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fn</span> <span class="o">=</span> <span class="n">fn</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">grad_fn</span></div>


<span class="k">class</span> <span class="nc">_Grad</span><span class="p">(</span><span class="n">GradOperation_</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A higher-order function which is used to generate the gradient function by position for the input function.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">get_by_list</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">sens_param</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">get_by_position</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _Grad.&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">get_by_position</span><span class="p">,</span> <span class="nb">bool</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;_Grad&#39;, the &#39;get_by_position&#39; should be bool, &quot;</span>
                            <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">get_by_position</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">get_by_list</span><span class="p">,</span> <span class="nb">bool</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;_Grad&#39;, the &#39;get_by_list&#39; should be bool, &quot;</span>
                            <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">get_by_list</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">sens_param</span><span class="p">,</span> <span class="nb">bool</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;_Grad&#39;, the &#39;sens_param&#39; should be bool, &quot;</span>
                            <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">sens_param</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">get_by_position</span> <span class="o">=</span> <span class="n">get_by_position</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">get_by_list</span> <span class="o">=</span> <span class="n">get_by_list</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sens_param</span> <span class="o">=</span> <span class="n">sens_param</span>
        <span class="n">GradOperation_</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;grad&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">,</span> <span class="n">get_by_list</span><span class="p">,</span> <span class="n">sens_param</span><span class="p">,</span> <span class="n">get_by_position</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">grad_fn</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fn</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">pynative_</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">grad_position</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="nf">_pynative_forward_run</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">grad</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">kwargs</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="n">grad_position</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot; Pynative forward run to build grad graph. &quot;&quot;&quot;</span>
        <span class="n">new_kwargs</span> <span class="o">=</span> <span class="n">kwargs</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">sens_param</span><span class="p">:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="s1">&#39;sens&#39;</span> <span class="ow">in</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                <span class="n">args</span> <span class="o">=</span> <span class="n">args</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">new_kwargs</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
                <span class="n">new_kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s1">&#39;sens&#39;</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">FunctionType</span><span class="p">):</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">check_run</span><span class="p">(</span><span class="n">grad</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">new_kwargs</span><span class="p">):</span>
                <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">set_grad_flag</span><span class="p">(</span><span class="kc">True</span><span class="p">)</span>
                <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">new_graph</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">new_kwargs</span><span class="p">)</span>
                <span class="n">output</span> <span class="o">=</span> <span class="n">fn</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">new_kwargs</span><span class="p">)</span>
                <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">end_graph</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">output</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">new_kwargs</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="c1"># Check if fn have run already</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">check_run</span><span class="p">(</span><span class="n">grad</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">new_kwargs</span><span class="p">):</span>
                <span class="n">fn</span><span class="o">.</span><span class="n">set_grad</span><span class="p">()</span>
                <span class="n">fn</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">new_kwargs</span><span class="p">)</span>
                <span class="n">fn</span><span class="o">.</span><span class="n">set_grad</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="n">weights</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">grad_position</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">grad_fn</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">fn</span> <span class="o">==</span> <span class="n">fn</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">grad_position</span> <span class="o">==</span> <span class="n">grad_position</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">grad_fn</span>
        <span class="n">grad_</span> <span class="o">=</span> <span class="n">_Grad</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">get_by_list</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">sens_param</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_by_position</span><span class="p">)</span>
        <span class="c1"># If calling Grad in GRAPH_MODE or calling Grad in ms_function, do grad in GRAPH_MODE</span>
        <span class="c1"># If calling Grad in pure PYNATIVE_MODE do grad in PYNATIVE_MODE</span>
        <span class="c1">#   In pure PYNATIVE_MODE the out layer after_grad just used to set pynative flag for inner GradOperation.</span>
        <span class="c1">#   In PYNATIVE_MODE calling Grad from ms_function, use the out layer after_grad do grad in GRAPH_MODE.</span>
        <span class="k">if</span> <span class="n">context</span><span class="o">.</span><span class="n">get_context</span><span class="p">(</span><span class="s2">&quot;mode&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="n">context</span><span class="o">.</span><span class="n">GRAPH_MODE</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_by_position</span><span class="p">:</span>
                <span class="nd">@ms_function</span>
                <span class="k">def</span> <span class="nf">after_grad</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">):</span>
                    <span class="k">return</span> <span class="n">grad_</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">weights</span><span class="p">,</span> <span class="n">grad_position</span><span class="p">)(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_by_list</span><span class="p">:</span>
                    <span class="nd">@ms_function</span>
                    <span class="k">def</span> <span class="nf">after_grad</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">):</span>
                        <span class="k">return</span> <span class="n">grad_</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">weights</span><span class="p">)(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="nd">@ms_function</span>
                    <span class="k">def</span> <span class="nf">after_grad</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">):</span>
                        <span class="k">return</span> <span class="n">grad_</span><span class="p">(</span><span class="n">fn</span><span class="p">)(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">pynative_</span><span class="p">:</span>
            <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">set_grad_position</span><span class="p">(</span><span class="n">grad_</span><span class="p">,</span> <span class="n">grad_position</span><span class="p">)</span>
            <span class="nd">@_wrap_func</span>
            <span class="k">def</span> <span class="nf">after_grad</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                <span class="k">if</span> <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">check_graph</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                    <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;Another grad step is running&quot;</span><span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">_pynative_forward_run</span><span class="p">(</span><span class="n">grad_</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">kwargs</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="n">grad_position</span><span class="p">)</span>
                <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">grad</span><span class="p">(</span><span class="n">grad_</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="n">weights</span><span class="p">,</span> <span class="n">grad_position</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                <span class="n">out</span> <span class="o">=</span> <span class="n">_pynative_executor</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                <span class="n">_pynative_executor</span><span class="o">.</span><span class="n">clear_grad</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                <span class="k">return</span> <span class="n">out</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">grad_</span><span class="o">.</span><span class="n">pynative_</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="c1"># after_grad of this branch can&#39;t use @ms_function, just directly call grad_</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_by_position</span><span class="p">:</span>
                <span class="k">def</span> <span class="nf">after_grad</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                    <span class="k">return</span> <span class="n">grad_</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">weights</span><span class="p">,</span> <span class="n">grad_position</span><span class="p">)(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_by_list</span><span class="p">:</span>
                    <span class="k">def</span> <span class="nf">after_grad</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                        <span class="k">return</span> <span class="n">grad_</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">weights</span><span class="p">)(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="k">def</span> <span class="nf">after_grad</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                        <span class="k">return</span> <span class="n">grad_</span><span class="p">(</span><span class="n">fn</span><span class="p">)(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">grad_fn</span> <span class="o">=</span> <span class="n">after_grad</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fn</span> <span class="o">=</span> <span class="n">fn</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">grad_position</span> <span class="o">=</span> <span class="n">grad_position</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">grad_fn</span>


<div class="viewcode-block" id="MultitypeFuncGraph"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.MultitypeFuncGraph.html#mindspore.ops.MultitypeFuncGraph">[docs]</a><span class="k">class</span> <span class="nc">MultitypeFuncGraph</span><span class="p">(</span><span class="n">MultitypeFuncGraph_</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Generates overloaded functions.</span>

<span class="sd">    MultitypeFuncGraph is a class used to generate overloaded functions, considering different types as inputs.</span>
<span class="sd">    Initialize an `MultitypeFuncGraph` object with name, and use `register` with input types as the decorator</span>
<span class="sd">    for the function to be registered. And the object can be called with different types of inputs,</span>
<span class="sd">    and work with `HyperMap` and `Map`.</span>

<span class="sd">    Args:</span>
<span class="sd">        name (str): Operator name.</span>
<span class="sd">        read_value (bool): If the registered function not need to set value on Parameter,</span>
<span class="sd">            and all inputs will pass by value, set `read_value` to True. Default: False.</span>

<span class="sd">    Raises:</span>
<span class="sd">        ValueError: If failed to find a matching function for the given arguments.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; # `add` is a metagraph object which will add two objects according to</span>
<span class="sd">        &gt;&gt;&gt; # input type using &quot;.register&quot; decorator.</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import ops</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import dtype as mstype</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.ops.composite import MultitypeFuncGraph</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; tensor_add = ops.Add()</span>
<span class="sd">        &gt;&gt;&gt; add = MultitypeFuncGraph(&#39;add&#39;)</span>
<span class="sd">        &gt;&gt;&gt; @add.register(&quot;Number&quot;, &quot;Number&quot;)</span>
<span class="sd">        ... def add_scala(x, y):</span>
<span class="sd">        ...     return x + y</span>
<span class="sd">        &gt;&gt;&gt; @add.register(&quot;Tensor&quot;, &quot;Tensor&quot;)</span>
<span class="sd">        ... def add_tensor(x, y):</span>
<span class="sd">        ...     return tensor_add(x, y)</span>
<span class="sd">        &gt;&gt;&gt; output = add(1, 2)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        3</span>
<span class="sd">        &gt;&gt;&gt; output = add(Tensor([0.1, 0.6, 1.2], dtype=mstype.float32), Tensor([0.1, 0.6, 1.2], dtype=mstype.float32))</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [0.2 1.2 2.4]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">,</span> <span class="n">read_value</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize MultitypeFuncGraph.&quot;&quot;&quot;</span>
        <span class="n">MultitypeFuncGraph_</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">entries</span> <span class="o">=</span> <span class="nb">list</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">read_value</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">set_signatures</span><span class="p">((</span>
                <span class="n">sig</span><span class="o">.</span><span class="n">make_sig</span><span class="p">(</span><span class="s1">&#39;args&#39;</span><span class="p">,</span> <span class="n">sig</span><span class="o">.</span><span class="n">sig_rw</span><span class="o">.</span><span class="n">RW_READ</span><span class="p">,</span> <span class="n">sig</span><span class="o">.</span><span class="n">sig_kind</span><span class="o">.</span><span class="n">KIND_VAR_POSITIONAL</span><span class="p">),))</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">entries</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">output</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">entries</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="mi">1</span><span class="p">](</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
            <span class="k">return</span> <span class="n">output</span>
        <span class="n">types</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="n">mstype</span><span class="o">.</span><span class="n">get_py_obj_dtype</span><span class="p">,</span> <span class="n">args</span><span class="p">))</span>
        <span class="k">for</span> <span class="n">sigs</span><span class="p">,</span> <span class="n">fn</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">entries</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">sigs</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">types</span><span class="p">):</span>
                <span class="k">continue</span>
            <span class="k">if</span> <span class="nb">any</span><span class="p">(</span><span class="ow">not</span> <span class="n">mstype</span><span class="o">.</span><span class="n">issubclass_</span><span class="p">(</span><span class="n">type_</span><span class="p">,</span> <span class="n">sig</span><span class="p">)</span> <span class="k">for</span> <span class="n">sig</span><span class="p">,</span> <span class="n">type_</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">sigs</span><span class="p">,</span> <span class="n">types</span><span class="p">)):</span>
                <span class="k">continue</span>
            <span class="n">output</span> <span class="o">=</span> <span class="n">fn</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
            <span class="k">return</span> <span class="n">output</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;MultitypeFuncGraph&#39;, cannot find fn match given args. Got (sigs, fn): </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">entries</span><span class="si">}</span><span class="s2">, &quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;and (dtype, args): </span><span class="si">{</span><span class="n">types</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>

<div class="viewcode-block" id="MultitypeFuncGraph.register"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.MultitypeFuncGraph.html#mindspore.ops.MultitypeFuncGraph.register">[docs]</a>    <span class="k">def</span> <span class="nf">register</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">type_names</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Register a function for the given type string.</span>

<span class="sd">        Args:</span>
<span class="sd">            type_names (Union[str, :class:`mindspore.dtype`]): Inputs type names or types list.</span>

<span class="sd">        Return:</span>
<span class="sd">            decorator, a decorator to register the function to run, when called under the</span>
<span class="sd">            types described in `type_names`.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">def</span> <span class="nf">deco</span><span class="p">(</span><span class="n">fn</span><span class="p">):</span>
            <span class="k">def</span> <span class="nf">convert_type</span><span class="p">(</span><span class="n">type_input</span><span class="p">):</span>
                <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">type_input</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
                    <span class="k">return</span> <span class="n">mstype</span><span class="o">.</span><span class="n">typing</span><span class="o">.</span><span class="n">str_to_type</span><span class="p">(</span><span class="n">type_input</span><span class="p">)</span>
                <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">type_input</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">Type</span><span class="p">):</span>
                    <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;MultitypeFuncGraph&#39;, register only support str or </span><span class="si">{</span><span class="n">mstype</span><span class="o">.</span><span class="n">Type</span><span class="si">}</span><span class="s2">, but got &quot;</span>
                                    <span class="sa">f</span><span class="s2">&quot;&#39;type_input&#39;: </span><span class="si">{</span><span class="n">type_input</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
                <span class="k">return</span> <span class="n">type_input</span>

            <span class="n">types</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="n">convert_type</span><span class="p">,</span> <span class="n">type_names</span><span class="p">))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">register_fn</span><span class="p">(</span><span class="n">type_names</span><span class="p">,</span> <span class="n">fn</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">entries</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">types</span><span class="p">,</span> <span class="n">fn</span><span class="p">))</span>
            <span class="k">return</span> <span class="n">fn</span>
        <span class="k">return</span> <span class="n">deco</span></div></div>


<div class="viewcode-block" id="HyperMap"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.HyperMap.html#mindspore.ops.HyperMap">[docs]</a><span class="k">class</span> <span class="nc">HyperMap</span><span class="p">(</span><span class="n">HyperMap_</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Hypermap will apply the set operation to input sequences.</span>

<span class="sd">    Apply the operations to every element of the sequence or nested sequence. Different</span>
<span class="sd">    from `Map`, the `HyperMap` supports to apply on nested structure.</span>

<span class="sd">    Args:</span>
<span class="sd">        ops (Union[MultitypeFuncGraph, None]): `ops` is the operation to apply. If `ops` is `None`,</span>
<span class="sd">            the operations should be put in the first input of the instance. Default is None.</span>
<span class="sd">        reverse (bool): The optimizer needs to be inverted in some scenarios to improve parallel performance,</span>
<span class="sd">          general users please ignore. `reverse` is the flag to decide if apply the operation reversely.</span>
<span class="sd">          Only supported in graph mode. Default is False.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **args** (Tuple[sequence]) - If `ops` is not `None`, all the inputs should be sequences with the same length.</span>
<span class="sd">          And each row of the sequences will be the inputs of the operation.</span>

<span class="sd">          If `ops` is `None`, the first input is the operation, and the others are inputs.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Sequence or nested sequence, the sequence of output after applying the function.</span>
<span class="sd">        e.g. `operation(args[0][i], args[1][i])`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `ops` is neither MultitypeFuncGraph nor None.</span>
<span class="sd">        TypeError: If `args` is not a Tuple.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, ops</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.ops.composite.base import MultitypeFuncGraph, HyperMap</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import dtype as mstype</span>
<span class="sd">        &gt;&gt;&gt; nest_tensor_list = ((Tensor(1, mstype.float32), Tensor(2, mstype.float32)),</span>
<span class="sd">        ...                     (Tensor(3, mstype.float32), Tensor(4, mstype.float32)))</span>
<span class="sd">        &gt;&gt;&gt; # square all the tensor in the nested list</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; square = MultitypeFuncGraph(&#39;square&#39;)</span>
<span class="sd">        &gt;&gt;&gt; @square.register(&quot;Tensor&quot;)</span>
<span class="sd">        ... def square_tensor(x):</span>
<span class="sd">        ...     return ops.square(x)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; common_map = HyperMap()</span>
<span class="sd">        &gt;&gt;&gt; output = common_map(square, nest_tensor_list)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        ((Tensor(shape=[], dtype=Float32, value= 1), Tensor(shape=[], dtype=Float32, value= 4)),</span>
<span class="sd">        (Tensor(shape=[], dtype=Float32, value= 9), Tensor(shape=[], dtype=Float32, value= 16)))</span>
<span class="sd">        &gt;&gt;&gt; square_map = HyperMap(square, False)</span>
<span class="sd">        &gt;&gt;&gt; output = square_map(nest_tensor_list)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        ((Tensor(shape=[], dtype=Float32, value= 1), Tensor(shape=[], dtype=Float32, value= 4)),</span>
<span class="sd">        (Tensor(shape=[], dtype=Float32, value= 9), Tensor(shape=[], dtype=Float32, value= 16)))</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ops</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">reverse</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize HyperMap.&quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ops</span> <span class="o">=</span> <span class="n">ops</span>
        <span class="k">if</span> <span class="n">ops</span><span class="p">:</span>
            <span class="n">HyperMap_</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">reverse</span><span class="p">,</span> <span class="n">ops</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">HyperMap_</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">reverse</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="n">func</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">ops</span>
        <span class="n">args_list</span> <span class="o">=</span> <span class="n">args</span>
        <span class="n">hypermap</span> <span class="o">=</span> <span class="bp">self</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">ops</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">func</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
            <span class="n">args_list</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="mi">1</span><span class="p">:]</span>
            <span class="n">hypermap</span> <span class="o">=</span> <span class="n">partial</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">func</span><span class="p">)</span>
        <span class="c1"># is leaf</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">args_list</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">)):</span>
            <span class="k">return</span> <span class="n">func</span><span class="p">(</span><span class="o">*</span><span class="n">args_list</span><span class="p">)</span>
        <span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="n">hypermap</span><span class="p">,</span> <span class="o">*</span><span class="n">args_list</span><span class="p">))</span></div>


<div class="viewcode-block" id="Map"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.Map.html#mindspore.ops.Map">[docs]</a><span class="k">class</span> <span class="nc">Map</span><span class="p">(</span><span class="n">Map_</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Map will apply the set operation on input sequences.</span>

<span class="sd">    Apply the operations to every element of the sequence.</span>

<span class="sd">    Args:</span>
<span class="sd">        ops (Union[MultitypeFuncGraph, None]): `ops` is the operation to apply. If `ops` is `None`,</span>
<span class="sd">            the operations should be put in the first input of the instance. Default: None</span>
<span class="sd">        reverse (bool): The optimizer needs to be inverted in some scenarios to improve parallel performance,</span>
<span class="sd">          general users please ignore. `Reverse` is the flag to decide if apply the operation reversely.</span>
<span class="sd">          Only supported in graph mode. Default is False.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **args** (Tuple[sequence]) - If `ops` is not `None`, all the inputs should be the same length sequences,</span>
<span class="sd">          and each row of the sequences. e.g. If the length of args is 2, and for `i` in length of each sequence</span>
<span class="sd">          `(args[0][i], args[1][i])` will be the input of the operation.</span>

<span class="sd">          If `ops` is `None`, the first input is the operation, and the other is inputs.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Sequence, the sequence of output after applying the function. e.g. `operation(args[0][i], args[1][i])`.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import dtype as mstype</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, ops</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.ops import MultitypeFuncGraph, Map</span>
<span class="sd">        &gt;&gt;&gt; tensor_list = (Tensor(1, mstype.float32), Tensor(2, mstype.float32), Tensor(3, mstype.float32))</span>
<span class="sd">        &gt;&gt;&gt; # square all the tensor in the list</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; square = MultitypeFuncGraph(&#39;square&#39;)</span>
<span class="sd">        &gt;&gt;&gt; @square.register(&quot;Tensor&quot;)</span>
<span class="sd">        ... def square_tensor(x):</span>
<span class="sd">        ...     return ops.square(x)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; common_map = Map()</span>
<span class="sd">        &gt;&gt;&gt; output = common_map(square, tensor_list)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        (Tensor(shape=[], dtype=Float32, value= 1), Tensor(shape=[], dtype=Float32, value= 4),</span>
<span class="sd">        Tensor(shape=[], dtype=Float32, value= 9))</span>
<span class="sd">        &gt;&gt;&gt; square_map = Map(square, False)</span>
<span class="sd">        &gt;&gt;&gt; output = square_map(tensor_list)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        (Tensor(shape=[], dtype=Float32, value= 1), Tensor(shape=[], dtype=Float32, value= 4),</span>
<span class="sd">        Tensor(shape=[], dtype=Float32, value= 9))</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ops</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">reverse</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize Map.&quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ops</span> <span class="o">=</span> <span class="n">ops</span>
        <span class="k">if</span> <span class="n">ops</span><span class="p">:</span>
            <span class="n">Map_</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">reverse</span><span class="p">,</span> <span class="n">ops</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">Map_</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">reverse</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="n">func</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">ops</span>
        <span class="n">args_list</span> <span class="o">=</span> <span class="n">args</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">ops</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">func</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
            <span class="n">args_list</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="mi">1</span><span class="p">:]</span>
        <span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="n">func</span><span class="p">,</span> <span class="o">*</span><span class="n">args_list</span><span class="p">))</span></div>


<span class="k">class</span> <span class="nc">Shard</span><span class="p">(</span><span class="n">Shard_</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Shard operation&quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize Shard.&quot;&quot;&quot;</span>
        <span class="n">Shard_</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;Shard&#39;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">shard_fn</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fn</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">in_axes</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">out_axes</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">device</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">level</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="n">in_axes</span><span class="p">,</span> <span class="n">out_axes</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">level</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">in_axes</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;Shard&#39;, the &#39;in_axes&#39; should be a tuple, but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">in_axes</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">out_axes</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;Shard&#39;, the &#39;out_axes&#39; should be a tuple, &quot;</span>
                            <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">out_axes</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">device</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;Shard&#39;, the &#39;device&#39; should be a string, &quot;</span>
                            <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">device</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;Shard&#39;, the &#39;level&#39; should be an integer, &quot;</span>
                            <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">level</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">shard_fn</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">fn</span> <span class="o">==</span> <span class="n">fn</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">in_axes</span> <span class="o">==</span> <span class="n">in_axes</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">out_axes</span> <span class="o">==</span> <span class="n">out_axes</span> <span class="ow">and</span> \
            <span class="bp">self</span><span class="o">.</span><span class="n">device</span> <span class="o">==</span> <span class="n">device</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">level</span> <span class="o">==</span> <span class="n">level</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">shard_fn</span>
        <span class="n">shard_</span> <span class="o">=</span> <span class="n">Shard</span><span class="p">()</span>

        <span class="nd">@ms_function</span>
        <span class="k">def</span> <span class="nf">after_shard</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">):</span>
            <span class="k">return</span> <span class="n">shard_</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">in_axes</span><span class="p">,</span> <span class="n">out_axes</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">level</span><span class="p">)(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">shard_fn</span> <span class="o">=</span> <span class="n">after_shard</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fn</span> <span class="o">=</span> <span class="n">fn</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">in_axes</span> <span class="o">=</span> <span class="n">in_axes</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">out_axes</span> <span class="o">=</span> <span class="n">out_axes</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">device</span> <span class="o">=</span> <span class="n">device</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">level</span> <span class="o">=</span> <span class="n">level</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">shard_fn</span>


<span class="k">class</span> <span class="nc">_ListAppend</span><span class="p">(</span><span class="n">ListAppend_</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A metafuncgraph class that append one element to list.</span>

<span class="sd">    Args:</span>
<span class="sd">        name (str): The name of the metafuncgraph object.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _ListAppend.&quot;&quot;&quot;</span>
        <span class="n">ListAppend_</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="k">pass</span>


<span class="n">_append</span> <span class="o">=</span> <span class="n">_ListAppend</span><span class="p">(</span><span class="s2">&quot;append&quot;</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">_ListInsert</span><span class="p">(</span><span class="n">ListInsert_</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A metafuncgraph class that insert one element to list.</span>

<span class="sd">    Args:</span>
<span class="sd">        name (str): The name of the metafuncgraph object.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _ListAppend.&quot;&quot;&quot;</span>
        <span class="n">ListInsert_</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="k">pass</span>


<span class="n">_insert</span> <span class="o">=</span> <span class="n">_ListInsert</span><span class="p">(</span><span class="s2">&quot;insert&quot;</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">_Tail</span><span class="p">(</span><span class="n">Tail_</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A metafuncgraph class that generates tail elements of the tuple.</span>

<span class="sd">    Args:</span>
<span class="sd">        name (str): The name of the metafuncgraph object.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _Tail.&quot;&quot;&quot;</span>
        <span class="n">Tail_</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="k">pass</span>


<span class="n">tail</span> <span class="o">=</span> <span class="n">_Tail</span><span class="p">(</span><span class="s1">&#39;tail&#39;</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">_ZipOperation</span><span class="p">(</span><span class="n">ZipOperation_</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Generates a tuple of zip iterations for inputs.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize _ZipOperation.&quot;&quot;&quot;</span>
        <span class="n">ZipOperation_</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="k">pass</span>


<span class="n">zip_operation</span> <span class="o">=</span> <span class="n">_ZipOperation</span><span class="p">(</span><span class="s1">&#39;zip_operation&#39;</span><span class="p">)</span>
<span class="sd">&quot;&quot;&quot;`zip_operation` will generate a tuple of zip iterations of inputs.&quot;&quot;&quot;</span>


<span class="n">env_get</span> <span class="o">=</span> <span class="n">MultitypeFuncGraph</span><span class="p">(</span><span class="s2">&quot;env_get&quot;</span><span class="p">)</span>


<span class="n">environ_get</span> <span class="o">=</span> <span class="n">Primitive</span><span class="p">(</span><span class="s1">&#39;EnvironGet&#39;</span><span class="p">)</span>
<span class="n">ref_to_embed</span> <span class="o">=</span> <span class="n">_grad_ops</span><span class="o">.</span><span class="n">RefToEmbed</span><span class="p">()</span>
<span class="n">zeros_like</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ZerosLike</span><span class="p">()</span>


<span class="nd">@env_get</span><span class="o">.</span><span class="n">register</span><span class="p">(</span><span class="s2">&quot;EnvType&quot;</span><span class="p">,</span> <span class="s2">&quot;Tensor&quot;</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">_tensor_env_get</span><span class="p">(</span><span class="n">env</span><span class="p">,</span> <span class="n">parameter</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Used to get env.&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">environ_get</span><span class="p">(</span><span class="n">env</span><span class="p">,</span> <span class="n">ref_to_embed</span><span class="p">(</span><span class="n">parameter</span><span class="p">),</span> <span class="n">zeros_like</span><span class="p">(</span><span class="n">parameter</span><span class="p">))</span>
</pre></div>

           </div>
           
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>
        &#169; Copyright 2021, MindSpore.

    </p>
  </div>
    
    
    
    Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    
    provided by <a href="https://readthedocs.org">Read the Docs</a>. 

</footer>
        </div>
      </div>

    </section>

  </div>
  

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
   

</body>
</html>